import sys
import os
import torch
from PyQt5 import QtWidgets, QtCore
from app_utils.get_model import get_hosipital_gpt
from app_utils.return_reply import hospitalgpt_reply

class HospitalGptChatApp(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()

        # 初始化模型和Tokenizer
        self.model = None
        self.tokenizer = None
        self.model_weights_dir = "weights/"  # 存放模型权重的文件夹

        self.init_ui()

    def init_ui(self):
        self.setWindowTitle("Hospital GPT")
        self.setGeometry(100, 100, 1100, 900)
        self.setStyleSheet("font-size: 12pt; font-family: Arial;")

        layout = QtWidgets.QVBoxLayout()

        # 聊天框
        self.chat_box = QtWidgets.QScrollArea(self)
        self.chat_box_widget = QtWidgets.QWidget(self.chat_box)
        self.chat_layout = QtWidgets.QVBoxLayout(self.chat_box_widget)
        self.chat_layout.setAlignment(QtCore.Qt.AlignTop)
        self.chat_box.setWidget(self.chat_box_widget)
        self.chat_box.setWidgetResizable(True)
        layout.addWidget(self.chat_box)

        # 模型选择
        self.model_label = QtWidgets.QLabel("选择模型权重：", self)
        self.model_select = QtWidgets.QComboBox(self)
        self.load_model_weights()  # 加载文件夹中的模型权重文件
        layout.addWidget(self.model_label)
        layout.addWidget(self.model_select)

        # 输入框
        self.input_box = QtWidgets.QLineEdit(self)
        self.input_box.setPlaceholderText("请输入您的症状描述...")
        self.input_box.setFixedHeight(60)
        self.input_box.setStyleSheet("padding: 10px; font-size: 30px; border-radius: 10px;")  # 设置样式
        layout.addWidget(self.input_box)

        # 发送按钮
        self.send_button = QtWidgets.QPushButton("发送", self)
        self.send_button.setFixedSize(100, 40)
        self.send_button.setStyleSheet("background-color: #4CAF50; color: white; border-radius: 5px; font-size: 14px;")
        self.send_button.clicked.connect(self.send_message)
        layout.addWidget(self.send_button)

        self.input_box.returnPressed.connect(self.send_message)

        self.setLayout(layout)

    def load_model_weights(self):

        if os.path.exists(self.model_weights_dir):
            weight_files = [f for f in os.listdir(self.model_weights_dir) if f.endswith('.pth') or f.endswith('.pt')]
            for file in weight_files:
                self.model_select.addItem(file.replace('.pth', '').replace('.pt', ''))  # 去掉扩展名
        else:
            self.chat_box.append("未找到模型权重文件夹！")

    def send_message(self):

        user_input = self.input_box.text()
        if user_input.strip():

            self.add_message_to_chat(f"<b>用户:</b> {user_input}", 'user')
            self.input_box.clear()

            response = self.get_model_response(user_input)
            for reply in response:
                self.add_message_to_chat(f"<b>问诊机器人:</b> {reply}", 'model')

    def add_message_to_chat(self, message, sender):

        message_widget = QtWidgets.QLabel(message, self)
        message_widget.setWordWrap(True)
        message_widget.setStyleSheet(f"""
            padding: 10px;
            background-color: {'#f0f0f0' if sender == 'user' else '#e1f5fe'};
            font-size: 25px;
            font-family: Arial;
            border-radius: 10px;
            border: 2px solid {'#ddd' if sender == 'user' else '#81d4fa'}; 
            margin-bottom: 10px;
        """)
        self.chat_layout.addWidget(message_widget)
        self.chat_box.verticalScrollBar().setValue(self.chat_box.verticalScrollBar().maximum())  # 自动滚动到最底部

    def load_model(self):

        model_name = self.model_select.currentText()
        self.model_name = self.model_select.currentText()

        if not model_name:
            self.add_message_to_chat("请选择一个有效的模型文件！", 'system')
            return

        model_path = os.path.join(self.model_weights_dir, model_name + ".pth")

        if not os.path.exists(model_path):
            self.add_message_to_chat(f"找不到模型文件: {model_path}", 'system')
            return

        # 加载词嵌入模型和匹配模型
        self.tokenizer, self.model_xlnet, self.model = get_hosipital_gpt(model_path)

    def get_model_response(self, user_input):

        if not self.model:
            self.load_model()
        elif self.model_name != self.model_select.currentText():
            print('选择新的模型，重新载入')
            self.load_model()

        medical_result_dialog, reply, medical_result, department_out, flag_call_120 = hospitalgpt_reply(user_input, self.tokenizer, self.model_xlnet, self.model)

        return medical_result_dialog, reply


def main():
    app = QtWidgets.QApplication(sys.argv)
    window = HospitalGptChatApp()
    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()
